# -*- coding: utf-8 -*-
import os
import sys
from collections import defaultdict

import torch
import re
import numpy as np

from model import XModel
from loader import XLoader


class XEvaluate():
    def __init__(self, loader:XLoader ,config):
        self.config = config
        self.data_test = loader.load_train_data("test")
        self.data_test = torch.LongTensor([data[0] for data in self.data_test]) #测试数据

        self.corpus_text,self.corpus_label = loader.load_evaluate_data("test")
        self.static_dict_true = self.count_data_ner(self.corpus_label) #统计样本标签
        # print("测试集样本统计：", self.static_dict_true)
        # xx = self.data_test[:,0:1]
        # print(self.data_test.shape)
        # 记录所有轮数的预测结果
        self.static_pred_result = {
            "LOCATION": {
                "pred_true_entity": [],
                "precision": [],
                "recall": [],
                "F1": []
            },
            "ORGANIZATION": {
                "pred_true_entity": [],
                "precision": [],
                "recall": [],
                "F1": []
            },
            "PERSON": {
                "pred_true_entity": [],
                "precision": [],
                "recall": [],
                "F1": []
            },"TIME": {
                "pred_true_entity": [],
                "precision": [],
                "recall": [],
                "F1": []
            }
        }
        

    def count_data_ner(self,labels):
        location_regx = "(04+)"
        organization_regx = "(15+)"
        person_regx = "(26+)"
        time_regx = "(37+)"

        static_dict = {
            "LOCATION": {
                "regx" : location_regx,
                "entity": [],
                "num": 0
            },
            "ORGANIZATION": {
                "regx": organization_regx,
                "entity": [],
                "num": 0
            },
            "PERSON": {
                "regx": person_regx,
                "entity": [],
                "num": 0
            },
            "TIME": {
                "regx": time_regx,
                "entity": [],
                "num": 0
            }
        }

        for index,label in enumerate(labels):
            sentence = self.corpus_text[index]
            label = label[:len(sentence)]
            label = "".join([str(t) for t in label])
            for key, static in static_dict.items():
                t_entity = []
                for location in re.finditer(static["regx"], label):
                    s, e = location.span()
                    t_entity.append(sentence[s:e])
                static["entity"].append(t_entity)
                static["num"] += len(t_entity)
        return static_dict


    def eval(self, model:XModel):
        model.eval()
        if torch.cuda.is_available():
            input_ids = self.data_test.cuda()

        with torch.no_grad():
            y_pred = model.forward(input_ids)
        if self.config.get("use_crf",False):
            y_pred_label = y_pred
            static_dict_pred = self.count_data_ner(np.array(y_pred_label)) #统计预测样本标签
        else:
            y_pred_label = torch.argmax(y_pred, dim=-1)
            static_dict_pred = self.count_data_ner(y_pred_label.cpu().detach().numpy()) #统计预测样本标签
        # print("预测集样本统计", static_dict_pred)



        # 正确率 = 识别出的正确实体数 / 识别出的实体数
        # 召回率 = 识别出的正确实体数 / 样本的实体数
        for key, static_dict in self.static_dict_true.items():
            print("实体类型：", key)
            print("测试集样本数：", self.static_dict_true[key]["num"])
            print("预测集样本数：", static_dict_pred[key]["num"])
            #预测正确的实体
            entity_test = self.static_dict_true[key]["entity"]
            entity_pred = static_dict_pred[key]["entity"]
            entity_pred_true_num = 0
            for e in entity_pred:
                if len(e) == 0:
                    continue
                if e in entity_test:
                    entity_pred_true_num +=1
            # print("预测正确的实体数：", entity_pred_true_num)
            precision = entity_pred_true_num / (static_dict_pred[key]["num"]+1e-5)
            # print("正确率：", precision)
            recall = entity_pred_true_num / (self.static_dict_true[key]["num"]+1e-5)
            # print("召回率：", recall)
            f1 = 2 * precision * recall / (precision + recall+1e-5)
            # print("F1值：", f1)
            print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
            self.static_pred_result[key]["pred_true_entity"].append(entity_pred_true_num)
            self.static_pred_result[key]["precision"].append('{:.2f}'.format(precision))
            self.static_pred_result[key]["recall"].append('{:.2f}'.format(recall))
            self.static_pred_result[key]["F1"].append('{:.2f}'.format(f1))
        print("-------------------------------------------------------------")
        for key, static_dict in self.static_pred_result.items():
            print("实体类型：", key)
            print("预测正确的实体总数：", self.static_pred_result[key]["pred_true_entity"])
            print("正确率：", self.static_pred_result[key]["precision"])
            print("召回率：", self.static_pred_result[key]["recall"])
            print("F1值：", self.static_pred_result[key]["F1"])
            print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")





#main
if __name__ == "__main__":
    config = {
        "train": "ner_data/train",
        "test": "ner_data/test",
        "dev": "ner_data/dev",
        "schema": "ner_data/schema.json",
        "vocab": "ner_data/vocab.json",
        "max_length": 128,
        "batch_size": 32
    }
    loader = XLoader(config)
    evaluate = XEvaluate(loader, config)